"""
    WRFReader

A comprehensive Julia module for reading WRF-ARW (Weather Research and Forecasting - Advanced Research WRF)
NetCDF background fields for GSI data assimilation.

This module handles:
- Reading WRF NetCDF output files (wrfout, wrfinput)
- Unstaggering variables from WRF's Arakawa C-grid
- Converting perturbation fields to total fields
- Extracting grid metadata and coordinates
- Converting to GSI BackgroundFields structure

# WRF Grid Staggering (Arakawa C-grid)

WRF uses an Arakawa C-grid where different variables are staggered:
- **U velocity**: Staggered in west-east direction (west_east_stag)
- **V velocity**: Staggered in south-north direction (south_north_stag)
- **Temperature, moisture**: At mass points (unstaggered)
- **Geopotential**: Staggered in vertical (bottom_top_stag)

# WRF Perturbation Fields

WRF stores some fields as perturbations from a base state:
- **Temperature**: T_total = T_pert + T_base (potential temperature)
- **Geopotential**: PH_total = PH + PHB
- **Pressure**: P_total = P + PB (optional, computed from other fields)

# Usage

```julia
using GSICoreAnalysis.DataIO.WRFReader

# Read WRF background file
bg_fields = read_wrf_netcdf("wrfinput_d01.mem0001")

# Access fields
println("Grid size: ", size(bg_fields.u))
println("Temperature range: ", extrema(bg_fields.t))
```
"""
module WRFReader

using NCDatasets
using Dates
using Printf
using LinearAlgebra

# Import parent module types
using ..DataIO: BackgroundFields
using ...GridOperations: RegionalGridConfig, RegionalGrid, initialize_grid

export read_wrf_netcdf

"""
    read_wrf_netcdf(filename::String) -> BackgroundFields

Read WRF-ARW NetCDF file and return BackgroundFields structure.

This function performs the following operations:
1. Opens the WRF NetCDF file
2. Reads grid dimensions and metadata
3. Reads atmospheric variables (U, V, T, QVAPOR, PSFC, PH, PHB)
4. Unstaggers U and V from WRF's Arakawa C-grid
5. Adds base state to perturbation fields
6. Converts potential temperature to temperature
7. Computes geopotential height
8. Creates grid configuration
9. Returns BackgroundFields structure

# Arguments
- `filename::String`: Path to WRF NetCDF file (wrfout_*, wrfinput_*)

# Returns
- `BackgroundFields`: Structure containing atmospheric fields and metadata

# WRF Variables Read
- `U`: West-East wind component (m/s) [Time, bottom_top, south_north, west_east_stag]
- `V`: South-North wind component (m/s) [Time, bottom_top, south_north_stag, west_east]
- `T`: Perturbation potential temperature (K) [Time, bottom_top, south_north, west_east]
- `QVAPOR`: Water vapor mixing ratio (kg/kg) [Time, bottom_top, south_north, west_east]
- `PSFC`: Surface pressure (Pa) [Time, south_north, west_east]
- `PH`: Perturbation geopotential (m²/s²) [Time, bottom_top_stag, south_north, west_east]
- `PHB`: Base state geopotential (m²/s²) [Time, bottom_top_stag, south_north, west_east]
- `P`: Perturbation pressure (Pa) [Time, bottom_top, south_north, west_east]
- `PB`: Base state pressure (Pa) [Time, bottom_top, south_north, west_east]
- `XLAT`: Latitude (degrees) [Time, south_north, west_east]
- `XLONG`: Longitude (degrees) [Time, south_north, west_east]

# Example
```julia
bg_fields = read_wrf_netcdf("/path/to/wrfinput_d01")
println("Dimensions: ", size(bg_fields.u))  # (nlat, nlon, nlev)
```
"""
function read_wrf_netcdf(filename::String)
    @info "Reading WRF NetCDF file: $filename"

    # Open NetCDF file
    NCDataset(filename, "r") do ds
        # Read dimensions
        time_idx = 1  # Use first time slice

        # Get dimension sizes
        # WRF dimensions: (west_east, south_north, bottom_top)
        # We need to transpose to Julia convention: (south_north, west_east, bottom_top)
        nx_stag = ds.dim["west_east_stag"]
        ny_stag = ds.dim["south_north_stag"]
        nx = ds.dim["west_east"]
        ny = ds.dim["south_north"]
        nz = ds.dim["bottom_top"]
        nz_stag = ds.dim["bottom_top_stag"]

        @info "WRF grid dimensions: nx=$nx, ny=$ny, nz=$nz (staggered: nx_stag=$nx_stag, ny_stag=$ny_stag, nz_stag=$nz_stag)"

        # Read staggered wind components
        # WRF NetCDF variable dimensions: (west_east_stag, south_north, bottom_top, Time) for U
        #                                  (west_east, south_north_stag, bottom_top, Time) for V
        U_stag = Array(ds["U"][:, :, :, time_idx])  # (west_east_stag, south_north, bottom_top)
        V_stag = Array(ds["V"][:, :, :, time_idx])  # (west_east, south_north_stag, bottom_top)

        # Already in correct order for unstaggering
        U_stag_reorder = U_stag  # (west_east_stag, south_north, bottom_top)
        V_stag_reorder = V_stag  # (west_east, south_north_stag, bottom_top)

        @info "Read staggered winds: U $(size(U_stag_reorder)), V $(size(V_stag_reorder))"

        # Unstagger winds to mass points
        u_unstag = unstagger_u(U_stag_reorder)  # (west_east, south_north, bottom_top)
        v_unstag = unstagger_v(V_stag_reorder)  # (west_east, south_north, bottom_top)

        # Transpose to Julia convention (south_north, west_east, bottom_top) -> (nlat, nlon, nlev)
        u = permutedims(u_unstag, (2, 1, 3))
        v = permutedims(v_unstag, (2, 1, 3))

        @info "Unstaggered winds: u $(size(u)), v $(size(v))"

        # Read mass point variables (already at mass points)
        # WRF NetCDF dimensions: (west_east, south_north, bottom_top, Time)
        T_pert = Array(ds["T"][:, :, :, time_idx])  # (west_east, south_north, bottom_top)
        QVAPOR = Array(ds["QVAPOR"][:, :, :, time_idx])  # (west_east, south_north, bottom_top)
        PSFC = Array(ds["PSFC"][:, :, time_idx])  # (west_east, south_north)

        # Transpose mass point variables to Julia convention (south_north, west_east, bottom_top)
        t_pert = permutedims(T_pert, (2, 1, 3))
        q = permutedims(QVAPOR, (2, 1, 3))
        q = @. max(q, 0.0)
        ps = permutedims(PSFC, (2, 1))

        @info "Read mass point fields: T $(size(t_pert)), Q $(size(q)), PS $(size(ps))"

        # Read geopotential (staggered in vertical)
        # WRF NetCDF dimensions: (west_east, south_north, bottom_top_stag, Time)
        PH = Array(ds["PH"][:, :, :, time_idx])  # (west_east, south_north, bottom_top_stag)
        PHB = Array(ds["PHB"][:, :, :, time_idx])  # (west_east, south_north, bottom_top_stag)

        # Add perturbation and base state
        PH_total = PH .+ PHB  # Total geopotential (m²/s²) in (west_east, south_north, bottom_top_stag)

        # Already in correct order for unstaggering
        PH_total_reorder = PH_total

        # Unstagger geopotential in vertical (average adjacent levels)
        PH_unstag = unstagger_vertical(PH_total_reorder)  # (west_east, south_north, bottom_top)

        # Convert geopotential to geopotential height: Z = PH/g
        g = 9.81  # m/s²
        Z_unstag = PH_unstag ./ g

        # Transpose geopotential height to Julia convention
        z = permutedims(Z_unstag, (2, 1, 3))

        @info "Computed geopotential height: z $(size(z))"

        # Read pressure for temperature conversion
        # WRF NetCDF dimensions: (west_east, south_north, bottom_top, Time)
        P_pert = Array(ds["P"][:, :, :, time_idx])  # (west_east, south_north, bottom_top)
        PB = Array(ds["PB"][:, :, :, time_idx])  # (west_east, south_north, bottom_top)
        P_total = P_pert .+ PB  # Total pressure (Pa)

        # Transpose to Julia convention (south_north, west_east, bottom_top)
        P_julia = permutedims(P_total, (2, 1, 3))

        # Convert potential temperature to temperature
        # T = θ * (P/P0)^(R/cp)
        # where P0 = 100000 Pa (reference pressure), R/cp = 0.2854
        P0 = 100000.0  # Pa
        Rd_cp = 287.04 / 1004.0  # R/cp for dry air ≈ 0.2854

        # θ_total = θ_pert + T0 (WRF base state potential temperature = 300K)
        T0 = 300.0  # K (WRF base state)
        theta = t_pert .+ T0

        # Convert to actual temperature (both in same (south_north, west_east, bottom_top) order)
        t = theta .* (P_julia ./ P0).^Rd_cp

        @info "Converted potential temperature to temperature: T range = $(extrema(t)) K"

        # Read coordinate variables
        # WRF NetCDF dimensions: (west_east, south_north, Time)
        XLAT = Array(ds["XLAT"][:, :, time_idx])  # (west_east, south_north)
        XLONG = Array(ds["XLONG"][:, :, time_idx])  # (west_east, south_north)

        # Transpose to Julia convention (south_north, west_east)
        lat = permutedims(XLAT, (2, 1))
        lon = permutedims(XLONG, (2, 1))

        # Read time
        times_str = ""
        try
            times_var = ds["Times"]
            if ndims(times_var) == 2
                chars = Array(times_var[:, time_idx])
                times_str = strip(String(vec(chars)))
            elseif ndims(times_var) == 1
                chars = Array(times_var)
                times_str = strip(String(vec(chars)))
            else
                times_str = strip(string(times_var[time_idx]))
            end
        catch err
            @warn "Failed to interpret WRF Times variable, using default analysis time" exception=(err, catch_backtrace())
        end

        time = isempty(times_str) ? DateTime(2021, 1, 1, 0, 0, 0) : parse_wrf_time(times_str)

        # Read WRF grid parameters for metadata (convert to Float64)
        map_proj = haskey(ds.attrib, "MAP_PROJ") ? Int(ds.attrib["MAP_PROJ"]) : 1
        truelat1 = haskey(ds.attrib, "TRUELAT1") ? Float64(ds.attrib["TRUELAT1"]) : 0.0
        truelat2 = haskey(ds.attrib, "TRUELAT2") ? Float64(ds.attrib["TRUELAT2"]) : 0.0
        stand_lon = haskey(ds.attrib, "STAND_LON") ? Float64(ds.attrib["STAND_LON"]) : 0.0
        dx_wrf = haskey(ds.attrib, "DX") ? Float64(ds.attrib["DX"]) : 0.0
        dy_wrf = haskey(ds.attrib, "DY") ? Float64(ds.attrib["DY"]) : 0.0
        cen_lat = haskey(ds.attrib, "CEN_LAT") ? Float64(ds.attrib["CEN_LAT"]) : 0.0
        cen_lon = haskey(ds.attrib, "CEN_LON") ? Float64(ds.attrib["CEN_LON"]) : 0.0

        # Map WRF projection types
        proj_name = get_wrf_projection_name(map_proj)

        @info "WRF projection: $proj_name, center: ($cen_lat, $cen_lon), dx=$dx_wrf m, dy=$dy_wrf m"

        # Create grid configuration
        # Note: nx, ny here are the unstaggered dimensions
        grid_config = RegionalGridConfig(
            nx, ny, nz;
            proj = proj_name,
            center_lon = cen_lon,
            center_lat = cen_lat,
            dx = dx_wrf,
            dy = dy_wrf,
            truelat1 = truelat1,
            truelat2 = truelat2
        )

        grid = initialize_grid(grid_config)

        # Create metadata dictionary
        metadata = Dict{String,Any}(
            "source_file" => filename,
            "model_type" => "WRF-ARW",
            "format" => "NetCDF",
            "grid_type" => proj_name,
            "map_proj" => map_proj,
            "dx" => dx_wrf,
            "dy" => dy_wrf,
            "center_lat" => cen_lat,
            "center_lon" => cen_lon,
            "stand_lon" => stand_lon,
            "truelat1" => truelat1,
            "truelat2" => truelat2,
            "dimensions" => (ny, nx, nz),
            "wrf_version" => haskey(ds.attrib, "TITLE") ? ds.attrib["TITLE"] : "Unknown",
            "latitude_range" => extrema(lat),
            "longitude_range" => extrema(lon)
        )

        @info "Created BackgroundFields with dimensions $(size(u))"
        @info "Field ranges: u=$(extrema(u)) m/s, v=$(extrema(v)) m/s, T=$(extrema(t)) K, q=$(extrema(q)) kg/kg, ps=$(extrema(ps)) Pa"

        # Return BackgroundFields structure
        return BackgroundFields(u, v, t, q, ps, z, grid, time; metadata=metadata)
    end
end

"""
    unstagger_u(U_stag::Array{T,3}) where T -> Array{Float64,3}

Unstagger U wind component from WRF's staggered grid.

U is staggered in the west-east (x) direction on WRF's Arakawa C-grid.
We average adjacent points to get values at mass points.

# Arguments
- `U_stag`: Staggered U component (west_east_stag, south_north, bottom_top)

# Returns
- Unstaggered U at mass points (west_east, south_north, bottom_top)
"""
function unstagger_u(U_stag::Array{T,3}) where T
    nx_stag, ny, nz = size(U_stag)
    nx = nx_stag - 1

    U = zeros(Float64, nx, ny, nz)

    for k in 1:nz
        for j in 1:ny
            for i in 1:nx
                U[i,j,k] = 0.5 * (U_stag[i,j,k] + U_stag[i+1,j,k])
            end
        end
    end

    return U
end

"""
    unstagger_v(V_stag::Array{T,3}) where T -> Array{Float64,3}

Unstagger V wind component from WRF's staggered grid.

V is staggered in the south-north (y) direction on WRF's Arakawa C-grid.
We average adjacent points to get values at mass points.

# Arguments
- `V_stag`: Staggered V component (west_east, south_north_stag, bottom_top)

# Returns
- Unstaggered V at mass points (west_east, south_north, bottom_top)
"""
function unstagger_v(V_stag::Array{T,3}) where T
    nx, ny_stag, nz = size(V_stag)
    ny = ny_stag - 1

    V = zeros(Float64, nx, ny, nz)

    for k in 1:nz
        for j in 1:ny
            for i in 1:nx
                V[i,j,k] = 0.5 * (V_stag[i,j,k] + V_stag[i,j+1,k])
            end
        end
    end

    return V
end

"""
    unstagger_vertical(field_stag::Array{T,3}) where T -> Array{Float64,3}

Unstagger field from vertically staggered grid.

Geopotential is staggered in the vertical (bottom_top_stag) direction.
We average adjacent levels to get values at mass levels.

# Arguments
- `field_stag`: Vertically staggered field (nx, ny, bottom_top_stag)

# Returns
- Unstaggered field at mass levels (nx, ny, bottom_top)
"""
function unstagger_vertical(field_stag::Array{T,3}) where T
    nx, ny, nz_stag = size(field_stag)
    nz = nz_stag - 1

    field = zeros(Float64, nx, ny, nz)

    for k in 1:nz
        for j in 1:ny
            for i in 1:nx
                field[i,j,k] = 0.5 * (field_stag[i,j,k] + field_stag[i,j,k+1])
            end
        end
    end

    return field
end

"""
    parse_wrf_time(time_str::String) -> DateTime

Parse WRF time string to DateTime object.

WRF time format: "YYYY-MM-DD_HH:MM:SS"

# Arguments
- `time_str`: WRF time string

# Returns
- `DateTime`: Parsed date and time
"""
function parse_wrf_time(time_str::AbstractString)
    # Remove any trailing/leading whitespace
    time_str = strip(time_str)

    # WRF format: "YYYY-MM-DD_HH:MM:SS"
    parts = split(time_str, '_')
    if length(parts) != 2
        @warn "Unexpected WRF time format: $time_str, using default time"
        return DateTime(2021, 1, 1, 0, 0, 0)
    end

    date_part = parts[1]
    time_part = parts[2]

    date_vals = parse.(Int, split(date_part, '-'))
    time_vals = parse.(Int, split(time_part, ':'))

    return DateTime(date_vals[1], date_vals[2], date_vals[3],
                    time_vals[1], time_vals[2], time_vals[3])
end

"""
    get_wrf_projection_name(map_proj::Integer) -> String

Convert WRF map projection code to projection name.

WRF map projection codes:
- 1: Lambert Conformal
- 2: Polar Stereographic
- 3: Mercator
- 6: Latitude-Longitude

# Arguments
- `map_proj`: WRF map projection code

# Returns
- `String`: Projection name
"""
function get_wrf_projection_name(map_proj::Integer)
    if map_proj == 1
        return "lambert_conformal"
    elseif map_proj == 2
        return "polar_stereographic"
    elseif map_proj == 3
        return "mercator"
    elseif map_proj == 6
        return "latlon"
    else
        @warn "Unknown WRF map projection code: $map_proj, defaulting to lambert_conformal"
        return "lambert_conformal"
    end
end

end # module WRFReader
